package com.taotao.job.schedule.lock;

import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author fanbeibei
 * @date 2021/3/7 19:16
 */
@Slf4j
@AllArgsConstructor
public class MysqlDistributedLock implements DistributedLock {
    private DataSource dataSource;

    private ConcurrentHashMap<String, ConnInfo> connInfoMap = new ConcurrentHashMap<>();

    @Getter
    @AllArgsConstructor
    private static class ConnInfo {
        private Connection conn = null;
        private Boolean connAutoCommit = null;
        private PreparedStatement preparedStatement = null;
    }


    @Override
    public boolean lock(String lockKey) {
        try {
            Connection conn = dataSource.getConnection();
            boolean connAutoCommit = conn.getAutoCommit();
            conn.setAutoCommit(false);
            PreparedStatement preparedStatement = conn.prepareStatement("SELECT lock_name FROM job_lock WHERE lock_name = ? FOR UPDATE");
            preparedStatement.setString(1, lockKey);
            preparedStatement.execute();

            connInfoMap.putIfAbsent(lockKey, new ConnInfo(conn, connAutoCommit, preparedStatement));
        } catch (Exception e) {
            release(lockKey);
            return false;
        }
        return true;
    }


    private void release(String lockKey) {
        ConnInfo connInfo = connInfoMap.get(lockKey);
        if (null == connInfo) {
            return;
        }

        connInfoMap.remove(lockKey);

        Connection conn = connInfo.getConn();
        try {
            conn.commit();
        } catch (SQLException e) {
            log.error("", e);
        }
        try {
            conn.setAutoCommit(connInfo.getConnAutoCommit());
        } catch (SQLException e) {
            log.error("", e);
        }
        try {
            conn.close();
        } catch (SQLException e) {
            log.error("", e);
        }

        // close PreparedStatement
        try {
            connInfo.getPreparedStatement().close();
        } catch (SQLException e) {
            log.error("", e);
        }
    }


    @Override
    public void releaseLock(String lockKey) {
        release(lockKey);
    }
}
